from typing import List, Dict, Any
from neo4j import GraphDatabase
from core.entities import Chemical, Hazard
from core.relations import EvidenceTriple, RelationType
from core.exceptions import ChemicalKGError

class Neo4jGraphService:
    def __init__(self, uri: str, auth: tuple):
        self._driver = GraphDatabase.driver(uri, auth=auth)

    def close(self):
        self._driver.close()

    def create_chemical(self, chemical: Chemical) -> str:
        with self._driver.session() as session:
            result = session.execute_write(
                self._create_chemical_node,
                chemical.model_dump()
            )
            return result["id"]

    @staticmethod
    def _create_chemical_node(tx, data: Dict[str, Any]):
        query = """
        CREATE (c:Chemical {
            id: $id,
            name: $name,
            en_name: $en_name,
            cas: $cas,
            un: $un,
            molecular_formula: $molecular_formula,
            physical_state: $physical_state,
            aliases: $aliases
        })
        RETURN c.id AS id
        """
        return tx.run(query, **data).single()

    def add_hazard(self, chem_id: str, hazard: Hazard) -> str:
        with self._driver.session() as session:
            return session.execute_write(
                self._link_hazard,
                chem_id,
                hazard.model_dump()
            )

    @staticmethod
    def _link_hazard(tx, chem_id: str, hazard_data: Dict[str, Any]):
        query = """
        MATCH (c:Chemical {id: $chem_id})
        MERGE (h:Hazard {code: $code})
        SET h += $hazard_data
        CREATE (c)-[r:HAS_HAZARD]->(h)
        RETURN h.code AS code
        """
        params = {"chem_id": chem_id, "code": hazard_data["code"], "hazard_data": hazard_data}
        return tx.run(query, **params).single()

    async def get_chemical_with_hazards(self, chem_id: str) -> Dict:
        """获取化学品及其关联危害"""
        query = """
        MATCH (c:Chemical {id: $id})-[:HAS_HAZARD]->(h:Hazard)
        RETURN c, COLLECT(h) AS hazards
        """
        return await self.session.run(query, id=chem_id)

    # 新增多模态关联
    async def add_evidence(self, triple: EvidenceTriple):
        """添加多模态证据"""
        rel_type = triple.relation.upper()
        query = f"""
        MATCH (head) WHERE head.id = $head_id
        MERGE (tail:{'Evidence' if triple.modality != 'text' else 'Literal'})
        CREATE (head)-[:{rel_type} $props]->(tail)
        """
        await self.session.run(query,
                               head_id=triple.head,
                               props=triple.model_dump()
                               )